# encoding: utf-8
from __future__ import print_function
import sys
sys.path.append('/opt/alps/lib')
sys.path.append('/share/opt/alps/lib')
import pyalps
import matplotlib.pyplot as plt
import pyalps.plot
import numpy as np
import time
import scipy.io
import glob,os,shutil


#----------------------1 Run the Task-------------------
a1 = time.time()
task_name = 'dmrg_no_interaction'
BHparms = []
sweeps=20
lanczos_tolerance=1e-12
BHparms.append({
    'LATTICE': "open chain lattice",
    'L': 30,
    'MODEL_LIBRARY': "no_interaction_model.xml",
    'MODEL' : "bosons_no_interaction",
    'E0' : 3.74349,
    'delta': 4.2829e-6,
    'N_total': 30,
    'CONSERVED_QUANTUMNUMBERS':'N',
    'Nmax':6,
    'MAXSTATES': 150,
    'SWEEPS': sweeps,
    'NUMBER_EIGENVALUES': 1,
    'LANCZOS_TOLERANCE':lanczos_tolerance
    })
# 删除上一次运行得到的文件
old_in_files = glob.glob('*%s*.in.*'%task_name)
old_out_files = glob.glob('*%s*.out.*'%task_name)
old_data_files = ['No_interaction_conv_data.npz']
for i in old_in_files + old_out_files + old_data_files:
    try:
        os.remove(i)  #不能删文件夹
    except OSError:
        shutil.rmtree(i)  #用来删文件夹
doc1 = open('1Run_print.txt','w')
print('report: This task has %d parallel tasks.' % (len(BHparms)))#, file=doc1)
print('task %s is Running...' % (task_name))#, file = doc1)
#doc1.close()

input_file = pyalps.writeInputFiles(task_name, BHparms)

pyalps.runApplication('mps_optim', input_file, writexml=True)

a2 = time.time()
#doc1 = open('1Run_print.txt','a')
print('rundmrg has spent %.1f seconds' % (a2 - a1))#, file = doc1)
#doc1.close()

#----------------------2 结果分析--------------------------
## 用spyder分析时请注释掉上面写input_file和runApplication的部分，这是我想到最好的办法了

result_files=pyalps.getResultFiles(prefix=task_name)
assert result_files!=[], 'Report: Error, no result file has been calculated out. The calculation has definitly failed.'
## 如果没有产生结果文件，马上报错，而不是隐藏着继续运行

iterations=pyalps.loadIterationMeasurements(result_files,what=['Energy','TruncatedWeight'])
Energy_vs_Iterations=pyalps.collectXY(iterations,x='iteration',y='Energy')
xdata=Energy_vs_Iterations[0].x #这array里面的元素居然是string，难怪不好好排序的
xdata=np.array([int(i) for i in xdata])
ydata=Energy_vs_Iterations[0].y #这array里面的元素是标准的float64

x_final=np.array([])
y_final=np.array([])
for nsweep in range(sweeps):
    mask=(xdata==nsweep)
    y_temp=ydata[mask]
    y_final=np.concatenate((y_final,y_temp))
    x_temp=np.linspace(nsweep,nsweep+1,len(y_temp)+1)
    x_temp=np.delete(x_temp,0)
    x_final=np.concatenate((x_final,x_temp))

low_standard=1e-11
y_final=y_final-y_final[-1]+2*low_standard   #-y_final[-1]+2e-7是观察能量的基准


## 保存在numpy的npz文件中，供别的Python程序使用
np.savez('No_interaction_conv_data.npz',sweeps=x_final,energy_difference=y_final)

## 保存为matlab的mat文件吧，这表示可以用熟悉的matlab处理
#scipy.io.savemat('MPSconvergency_history_noInteraction.mat', {
#    'MPS_sweeps': x_final,
#    'Energy_difference': y_final,
#    'low_standard':low_standard,
#    'model_name':'NoInteractionModel'
#})

## 画图先不谈了
fig=plt.figure()
(fig,ax)=plt.subplots(figsize=(4,3),dpi=300)
ax.plot(x_final,y_final,'b')#,label='|Magneti|')
#ax.plot(x2,y2,'--',color='purple',label='Energy')
ax.legend(loc=4)
ax.set_xlabel('MPS sweeps')
ax.set_ylabel('Energy Difference')
ax.set_yscale('log')
#ax.set_ylim([low_standard,1e-4])
ax.grid()
ax.set_title('No interaction y 1e-12 limit')
fig.savefig(task_name + '_sweeps_vs_energy.pdf')
 